"""App entry: config, device, video loop, YOLO, ROI, alerts."""
from __future__ import annotations
import argparse
import time
from pathlib import Path
import os
import cv2  # type: ignore

from utils.config import load_config
from utils.device import select_device
from utils.logging import get_logger
from inference.yolo_detector import YoloDetector
from logic.roi_manager import load_roi_json, save_roi_json, interactive_define_roi, ROI
from utils.geometry import point_in_polygon, approx_distance_polygon_to_bbox, center_distance
from logic.alert_engine import ViolationState, update_state
from zny_io.video_reader import VideoReader
from zny_io.video_writer import VideoWriter
from zny_io.overlay import put_text, draw_bbox, draw_polygon
from zny_io.audio import play_wav


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--cfg", type=str, default="configs/config.yaml")
    p.add_argument("--source", type=str, default=None)
    return p.parse_args()


def _resolve_source(src: str | int | None):
    if src is None:
        return 0  # default camera
    # numeric camera id
    try:
        if isinstance(src, str) and src.isdigit():
            return int(src)
    except Exception:
        pass
    p = Path(str(src))
    if p.exists():
        return str(p)
    # fallback camera
    return 0


def main():
    args = parse_args()
    cfg = load_config(args.cfg)
    if args.source:
        cfg["app"]["source"] = args.source

    logger = get_logger(level=cfg.get("logging", {}).get("level", "INFO"),
                        file_path=cfg.get("logging", {}).get("file", None))
    logger.info("Starting ZebraNoYield")

    device, _ = select_device(cfg.get("device", {}).get("policy", "auto"),
                              prefer_fp16=cfg.get("device", {}).get("prefer_fp16", True))
    logger.info(f"Device selected: {device}")

    # Prepare model
    weights = cfg.get("model", {}).get("weights", "weights/yolov8n.pt")
    imgsz = int(cfg.get("model", {}).get("imgsz", 640))
    conf = float(cfg.get("model", {}).get("conf", 0.35))
    iou = float(cfg.get("model", {}).get("iou", 0.45))
    classes = cfg.get("model", {}).get("classes", ["person", "car", "bus", "truck", "motorcycle", "bicycle"])  # not directly used
    person_label = "person"
    vehicle_labels = [c for c in classes if c != person_label]

    detector = YoloDetector(weights=weights, device=device)

    # Prepare video
    src = _resolve_source(cfg.get("app", {}).get("source"))
    try:
        reader = VideoReader(src)
    except Exception as e:
        logger.error(f"Cannot open source {src}: {e}")
        return
    fps = reader.cap.get(cv2.CAP_PROP_FPS) or 30
    width = int(reader.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(reader.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # ROI
    roi_path = cfg.get("roi", {}).get("file", "configs/crosswalk.roi.json")
    roi: ROI | None = load_roi_json(roi_path)
    if roi is None:
        logger.info("ROI file not found. Press and click to define ROI polygon, 's'/Enter to save, 'q' to cancel.")
        ok, frame0 = reader.read()
        if not ok:
            logger.error("Failed to read first frame for ROI definition.")
            return
        roi = interactive_define_roi(frame0)
        if roi is None:
            logger.warning("ROI not defined. Using default horizontal band.")
            # default ROI: lower middle band
            pts = [(int(0.15 * width), int(0.65 * height)), (int(0.85 * width), int(0.65 * height)),
                   (int(0.85 * width), int(0.85 * height)), (int(0.15 * width), int(0.85 * height))]
            roi = ROI.from_pixels(pts, width, height)
        else:
            save_roi_json(roi_path, roi, (width, height))
            logger.info(f"ROI saved to {roi_path}")

    # Writer (optional)
    out_path = cfg.get("app", {}).get("output_video")
    writer = None
    if out_path:
        Path(out_path).parent.mkdir(parents=True, exist_ok=True)
        try:
            writer = VideoWriter(out_path, fps=fps or 30, size=(width, height))
            logger.info(f"Writing annotated video to {out_path}")
        except Exception as e:
            logger.warning(f"Open writer failed: {e}")

    # Logic thresholds
    logic_cfg = cfg.get("logic", {})
    D_roi_px = int(logic_cfg.get("D_roi_px", 60))
    D_pp_px = int(logic_cfg.get("D_pp_px", 80))
    vs = ViolationState()

    window = "ZebraNoYield"
    cv2.namedWindow(window)
    paused = False
    last_alert_ts = 0.0
    audio_file = cfg.get("audio", {}).get("file", "assets/sounds/alert.wav")
    screenshots_dir = Path("outputs/screenshots")
    last_raw = None
    last_disp = None

    while True:
        if not paused:
            ok, frame = reader.read()
            if not ok:
                logger.info("End of stream or read failure.")
                break

            # Inference
            persons, vehicles = detector.detect_persons_vehicles(
                frame, imgsz=imgsz, conf=conf, iou=iou, person_label=person_label, vehicle_labels=vehicle_labels
            )

            # ROI & distances
            roi_pts = roi.to_pixels(width, height) if roi else []
            persons_in_roi = [p for p in persons if point_in_polygon(p.center, roi_pts)] if roi else []
            vehicle_near_roi = any(approx_distance_polygon_to_bbox(roi_pts, v.bbox_xyxy) <= D_roi_px for v in vehicles) if roi else False
            vehicle_near_person = any(center_distance(p.center, v.center) <= D_pp_px for p in persons_in_roi for v in vehicles)

            alert, status = update_state(bool(persons_in_roi), bool(vehicle_near_roi or vehicle_near_person), vs, logic_cfg)

            # OSD
            last_raw = frame.copy()
            if roi:
                draw_polygon(frame, roi_pts, color=(0, 255, 255))
            for p in persons:
                draw_bbox(frame, p.bbox_xyxy, (0, 255, 0), f"P {p.conf:.2f}")
            for v in vehicles:
                draw_bbox(frame, v.bbox_xyxy, (255, 0, 0), f"V {v.conf:.2f}")
            put_text(frame, f"Status: {status}", (10, 24), (0, 255, 0) if status != "Alarming" else (0, 0, 255))
            put_text(frame, f"P_in_ROI:{len(persons_in_roi)} V:{len(vehicles)}", (10, 50), (200, 200, 200), 0.6)

            if alert is not None:
                put_text(frame, "NO YIELD ALERT!", (10, 80), (0, 0, 255), 0.9)
                # audio (best-effort, non-blocking)
                if audio_file and os.path.exists(audio_file):
                    try:
                        play_wav(audio_file)
                    except Exception:
                        pass

            if writer:
                writer.write(frame)

            last_disp = frame
            cv2.imshow(window, last_disp)
        # keys
        k = cv2.waitKey(1) & 0xFF
        if k == ord('q') or k == 27:
            break
        elif k == ord('p') or k == 32:  # pause or space
            paused = not paused
        elif k == ord('g'):
            if last_disp is not None:
                screenshots_dir.mkdir(parents=True, exist_ok=True)
                ts = int(time.time() * 1000)
                outp = screenshots_dir / f"shot_{ts}.jpg"
                try:
                    cv2.imwrite(str(outp), last_disp)
                    logger.info(f"Saved screenshot: {outp}")
                except Exception as e:
                    logger.warning(f"Save screenshot failed: {e}")
        elif k == ord('e') or k == ord('E'):
            # Try to fetch a frame for editing
            src_for_edit = None
            if last_raw is not None:
                src_for_edit = last_raw
            elif last_disp is not None:
                src_for_edit = last_disp
            else:
                ok2, fr2 = reader.read()
                if ok2:
                    src_for_edit = fr2
            if src_for_edit is not None:
                logger.info("Entering ROI editor (left add, right undo, r reset, s/Enter save, q cancel)")
                new_roi = interactive_define_roi(src_for_edit)
                if new_roi is not None:
                    roi = new_roi
                    save_roi_json(roi_path, roi, (width, height))
                    logger.info(f"ROI updated and saved: {roi_path}")
                else:
                    logger.info("ROI edit canceled")

    if writer:
        writer.release()
    reader.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
